#include "lit/elf/elfFile.h"
#include "lit/elf/linker.h"
#include "lit/report/report.h"

void Linker::allocAddr() {
    u32 curAddr = BASE_ADDR;
    u32 curOff = sizeof(Elf32_Ehdr) + sizeof(Elf32_Phdr) * segNames.size();
    for (u32 i = 0; i < segNames.size(); ++i) {
        segLists[segNames[i]]->allocAddr(segNames[i], curAddr, curOff);
    }
}

void Linker::collectInfo() {
    for (u32 i = 0; i < elfs.size(); ++i) {
        // 遍历所有的elf文件
        auto elf = elfs[i];
        // 记录段表信息
        for (u32 j = 0; j < segNames.size(); ++j) {
            // 遍历每个elf文件中需要处理的段（.text .data)
            if (elf->shdrTable.find(segNames[j]) != elf->shdrTable.end()) {
                // segLists记录 段名：拥有该段的所有elf文件
                segLists[segNames[j]]->ownerList.emplace_back(elf);
            }
        }
        // 遍历所有elf文件的符号表
        for (auto symIter = elf->symbolTable.begin(); symIter != elf->symbolTable.end(); ++symIter) {
            // 只记录全局符号
            if (ELF32_ST_BIND(symIter->second->st_info) == STB_GLOBAL) {
                auto symLink = std::make_shared<SymLink>();
                symLink->name = symIter->first;
                // 如果符号所在段是未定义的，则该符号在该elf文件中是外部符号
                if (symIter->second->st_shndx == STN_UNDEF) {
                    // 设置外部符号的接收者为当前elf文件，提供者为空
                    symLink->recv = elf;
                    symLink->prov = nullptr;
                    symLinks.emplace_back(symLink);
                } else {
                    // 如果符号所在段是已定义的，则该符号的提供者是当前elf文件，接收者未知
                    symLink->recv = nullptr;
                    symLink->prov = elf;
                    symDef.emplace_back(symLink);
                }
            }
        }
    }
}

void Linker::symValid() {
    // 程序入口的elf文件初始化为空
    startOwner = nullptr;
    // 遍历所有定义的符号
    for (u32 i = 0; i < symDef.size(); ++i) {
        // 如果符号名为start，则程序入口的elf所有者为start符号的提供者
        if (symDef[i]->name == START) {
            startOwner = symDef[i]->prov;
        }
        // 遍历当前符号之后的其他符号，如果之后有符号名称与当前符号相同，则符号被多次定义，报告错误
        for (u32 j = i + 1; j < symDef.size(); ++j) {
            if (symDef[i]->name == symDef[j]->name) {
                fatalMacro(("符号重定义，符号名为：" + symDef[i]->name).c_str());
            }
        }
    }

    // 如果遍历完所有的符号，并未发现入口符号的定义，报告错误
    if (startOwner == nullptr) {
        fatalMacro("找不到程序的入口点");
    }

    // 遍历所有外部引用的符号
    for (u32 i = 0; i < symLinks.size(); ++i) {
        // 如果未在定义符号集合中找到该引用符号的定义，报告错误“未定义的引用”
        for (u32 j = 0; j < symDef.size(); ++j) {
            // 如果找到该引用符号的定义，更新该引用符号的提供者
            if (symLinks[i]->name == symDef[j]->name) {
                symLinks[i]->prov = symDef[j]->prov;
            }
        }

        if (symLinks[i]->prov == nullptr) {
            fatalMacro(("未定义的符号 ： " + symLinks[i]->name).c_str());
        }
    }
}

void Linker::symParser() {
    // 遍历所有的elf文件
    for (u32 i = 0; i < elfs.size(); ++i) {
        auto elf = elfs[i];
        // 遍历所有的符号
        for (auto symIter = elf->symbolTable.begin(); symIter != elf->symbolTable.end(); ++symIter) {
            auto sym = symIter->second;
            // 如果为本地符号，将本地符号的值修改为分配的虚拟地址
            if (sym->st_shndx != STN_UNDEF) {
                std::string segName = elf->shdrNames[sym->st_shndx];
                sym->st_value += elf->shdrTable[segName]->sh_addr;
            }
        }
    }
    // 遍历所有的外部符号，将符号的接收者的值设置为符号提供者的值
    for (u32 i = 0; i < symLinks.size(); ++i) {
        std::string name = symLinks[i]->name;
        auto provSym = symLinks[i]->prov->symbolTable[name];
        auto recvSym = symLinks[i]->recv->symbolTable[name];
        recvSym->st_value = provSym->st_value;
    }
}

void Linker::relocate() {
    // 遍历所有的elf文件
    for (int i = 0; i < elfs.size(); ++i) {
        // 遍历每一个elf文件的重定位表
        std::vector<std::shared_ptr<RelocationItem>> tab = elfs[i]->relTable;
        for (int j = 0; j < tab.size(); ++j) {
            // 遍历重定位项
            // 重定位符号信息
            std::string relName = tab[j]->m_relName;
            auto sym = elfs[i]->symbolTable[relName];
            // 重定位段
            std::string segName = tab[j]->m_segName;
            auto seg = elfs[i]->shdrTable[segName];
            // 重定位符号地址
            u32 symAddr = sym->st_value;
            // 重定位位置
            u32 offset = tab[j]->m_rel->r_offset;
            u32 relAddr = seg->sh_addr + offset;
            // 重定位类型
            u8 type = ELF32_R_TYPE(tab[j]->m_rel->r_info);
            //重定位操作
            segLists[tab[j]->m_segName]->relocAddr(relAddr, type, symAddr);
        }
    }
}

void Linker::addElf(const char *elfFileName) {
    if (elfFileName == nullptr) {
        return;
    }
    auto elf = std::make_shared<ElfFile>();
    elf->readElf(elfFileName);
    elfs.emplace_back(elf);
}

void Linker::assemExe() {
    exe = std::make_shared<ElfFile>();
    exe->assemObj(this);
}

void Linker::exportElf(const char *fileName) {
    exe->writeElf(this, fileName);
}
